/*
 * SRMI (an easy-to-use zero-code remoting library for Java)
 * Copyright (C) 2007  Mohammad Nabil Hussein
 * mohammad (dot) nabil (dot) h (at) gmail (dot) com
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 */
package SRMI.Tests;

import java.io.IOException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtPrimitiveType;
import javassist.bytecode.Bytecode;
import javassist.compiler.Javac;
import org.apache.bcel.Constants;
import org.apache.bcel.classfile.JavaClass;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.BranchInstruction;
import org.apache.bcel.generic.INVOKESPECIAL;
import org.apache.bcel.generic.Instruction;
import org.apache.bcel.generic.InstructionConstants;
import org.apache.bcel.generic.InstructionFactory;
import org.apache.bcel.generic.InstructionHandle;
import org.apache.bcel.generic.InstructionList;
import org.apache.bcel.generic.InstructionTargeter;
import org.apache.bcel.generic.MethodGen;
import org.apache.bcel.generic.RETURN;
import org.apache.bcel.generic.TargetLostException;
import org.apache.bcel.generic.Type;
import org.apache.bcel.util.ByteSequence;

public class DefaultCtorNeutralizer extends BasicInstrumenter
{
    public DefaultCtorNeutralizer(JavaClass clazz)
    {
        super(clazz);
    }

    public DefaultCtorNeutralizer(String path) throws IOException
    {
        super(path);
    }

    public DefaultCtorNeutralizer(byte[] clazz, String name) throws IOException
    {
        super(clazz, name);
    }

    public DefaultCtorNeutralizer(Class<?> clazz) throws IOException
    {
        super(clazz);
    }

    private Method getArglessMethod(String methodName)
    {
        for (Method m : jclas.getMethods())
        {
            if (m.getName().equals(methodName) && m.getArgumentTypes().length == 0)
            {
                return m;
            }
        }
        return null;
    }

    private void removeFirstInstruction(InstructionList il)
    {
        InstructionHandle toRem = il.findHandle(il.getInstructionPositions()[0]);
        InstructionHandle newTarget = il.findHandle(il.getInstructionPositions()[1]);

        if (toRem.getTargeters() != null)
        {
            for (InstructionTargeter targeter : toRem.getTargeters())
            {
                targeter.updateTarget(toRem, newTarget);
            }
        }

        try
        {
            il.delete(toRem);
        }
        catch (TargetLostException ex)
        {
            ex.printStackTrace();
            System.exit(1);
        }

        il.setPositions();
    }

    private void renameDefaultCtor()
    {
        InstructionList il = new MethodGen(getArglessMethod("<init>"), jclas.getClassName(), cgen.getConstantPool()).getInstructionList();

        // remove super call, cuz already called by stub
        removeFirstInstruction(il); // remove aload_0
        removeFirstInstruction(il); // remove invokespecial ${super}.<init>

        MethodGen methodGen = new MethodGen(
                Constants.ACC_PUBLIC | Constants.ACC_SYNTHETIC,
                Type.VOID,
                Type.NO_ARGS,
                null /*param names*/,
                oldCtorName,
                cgen.getClassName(),
                il,
                cgen.getConstantPool());

        methodGen.stripAttributes(false);
        methodGen.setMaxStack();
        methodGen.setMaxLocals();
        cgen.removeMethod(getArglessMethod("<init>"));
        cgen.replaceMethod(getArglessMethod(oldCtorName), methodGen.getMethod());
        il.dispose();
    }

    void createNuetralCtor() {
        InstructionList il = new InstructionList();
        il.append(_factory.createPrintln(jclas.getClassName() + "'s aritificial ctor called"));
        il.append(InstructionConstants.RETURN);
        MethodGen mg = new MethodGen(
                Constants.ACC_PUBLIC | Constants.ACC_SYNTHETIC,
                Type.VOID,
                Type.NO_ARGS,
                null,
                neutralCtorName,
                cgen.getClassName(),
                il,
                cgen.getConstantPool());
        mg.stripAttributes(false);
        mg.setMaxStack();
        mg.setMaxLocals();
        cgen.replaceMethod(getArglessMethod(neutralCtorName), mg.getMethod());
        il.dispose();
    }

    private void createStubCtor() throws ClassNotFoundException
    {
        InstructionList il = new InstructionList();
        InstructionList body = null;

        boolean useJavassit = false;

        if (useJavassit)
        {
            try
            {
                CtClass cc = ClassPool.getDefault().get(jclas.getClassName());
                Javac jv = new Javac(cc);
                Bytecode b = jv.compileBody(
                        new javassist.CtMethod(CtPrimitiveType.voidType,
                                               "<init>", new CtClass[]{}, cc), "{System.out.println(\"gen\");}");
                // MUST do something to convert the indices from that class pool
                // to our class pool.
                // BUT i can, with some work, know the indices
                // and the strings, but how do i know which string
                // is a method ref or a class ref or a field ref ???
                //cc.getClassFile().getConstPool().
                body = new InstructionList(b.get());
                if (Config.verbose)
                {
                    ByteSequence bytes = new ByteSequence(b.get());
                    while (bytes.available() > 0)
                    {
                        Instruction i = Instruction.readInstruction(bytes);
                        System.out.println(i);
                    }
                }
            }
            catch (Exception ex)
            {
                ex.printStackTrace();
            }
        }

        il.append(InstructionConstants.THIS);
        il.append(new INVOKESPECIAL(cgen.getConstantPool().addMethodref(
                                    jclas.getSuperclassName(), "<init>", "()V")));

        if (useJavassit)
        {
            il.append(body);
        }
        else
        {
            il.append(
                    _factory.createFieldAccess(
                    DeciderClassName, DecisionFieldName, Type.BOOLEAN, Constants.GETSTATIC));
            BranchInstruction ifeq_7 =
                    InstructionFactory.createBranchInstruction(Constants.IFEQ, null);
            il.append(ifeq_7);

            il.append(InstructionConstants.THIS);
            il.append(new INVOKESPECIAL(cgen.getConstantPool().addMethodref(jclas.getClassName(), neutralCtorName, "()V")));

            BranchInstruction goto_14 =
                    InstructionFactory.createBranchInstruction(Constants.GOTO, null);
            il.append(goto_14);

            InstructionHandle ih_17 =
                    il.append(InstructionConstants.THIS);
            il.append(new INVOKESPECIAL(cgen.getConstantPool().addMethodref(jclas.getClassName(), oldCtorName, "()V")));

            InstructionHandle ih_21 =
                    il.append(InstructionConstants.RETURN);

            ifeq_7.setTarget(ih_17);
            goto_14.setTarget(ih_21);
        }

        MethodGen mg = new MethodGen(Constants.ACC_PUBLIC,
                                     Type.VOID,
                                     Type.NO_ARGS,
                                     null,
                                     "<init>",
                                     cgen.getClassName(), il, cgen.getConstantPool());
        mg.stripAttributes(false);
        mg.setMaxStack();
        mg.setMaxLocals();
        cgen.addMethod(mg.getMethod());
        il.dispose();
    }

    public byte[] instrument() throws Exception
    {

        // must do the check using BCEL, cuz if i used java reflection the class will be loaded
        // which is what we don't want to do before we load the class
        // we can't even use another class loader, cuz reflection uses the caller's class loader
        if (alreadyPatched())
        {
            if (Config.verbose)
            {
                System.out.println("Class " + jclas.getClassName() + " already patched");
            }
            return jclas.getBytes();
        }

        createStubCtor();

        Method oldCtor = getArglessMethod("<init>");
        if (oldCtor == null)
        {
            if (Config.verbose)
            {
                System.out.println("!" + jclas.getClassName() + " doesn't have a ctor to patch");
            }
            return jclas.getBytes();
        }

        renameDefaultCtor();

        createNuetralCtor();

        //flushUpdatedClass();

        if (Config.verbose)
        {
            System.out.println("+" + jclas.getClassName() + " patched");
        }

        verify(jclas.getClassName());

        return cgen.getJavaClass().getBytes();
    }

//    void flushUpdatedClass() throws Exception
//    {
//        FileOutputStream fs = new FileOutputStream("./build/classes/" + jclas.getClassName().replace('.', '_') + ".class");
//        cgen.getJavaClass().dump(fs);
//        fs.close();
//    }

//    void useJavassist(String qualifiedClassName) throws Exception {
//        CtClass clazz = ClassPool.getDefault().get(qualifiedClassName);
//        for (CtConstructor ctor : clazz.getConstructors()) {
//            ctor.setBody("{ System.out.println(\"artificialism\"); }");
//        }
//        clazz.writeFile("./build/classes");
//        verify(clazz.getName());
//    }
    
    boolean alreadyPatched()
    {
        for (Method m : jclas.getMethods())
        {
            if (m.getName().equals(oldCtorName))
            {
                return true;
            }
        }
        return false;
    }

    public byte[] prepare()
    {
        if (alreadyPatched())
        {
            if (Config.verbose)
            {
                System.out.println("Class " + jclas.getClassName() + " already prepared");
            }
            return jclas.getBytes();
        }

        MethodGen mgOld = new MethodGen(Constants.ACC_PUBLIC | Constants.ACC_SYNTHETIC, Type.VOID, Type.NO_ARGS, null, oldCtorName, cgen.getClassName(), new InstructionList(new RETURN()), cgen.getConstantPool());
        MethodGen mgNew = new MethodGen(Constants.ACC_PUBLIC | Constants.ACC_SYNTHETIC, Type.VOID, Type.NO_ARGS, null, neutralCtorName, cgen.getClassName(), new InstructionList(new RETURN()), cgen.getConstantPool());
        cgen.addMethod(mgOld.getMethod());
        cgen.addMethod(mgNew.getMethod());

        return cgen.getJavaClass().getBytes();
    }
    public static boolean good = true;
    public final static String oldCtorName = "ctor$oldCtor";
    public final static String neutralCtorName = "ctor$newCtor";
    public final static String DeciderClassName = "SRMI.Tests.DefaultCtorNeutralizer";
    public final static String DecisionFieldName = "good";

    public static void main(String[] argv) throws Exception
    {
        System.out.println("Evil scenario:");
        new MyServer();
        
        useInstrumenter(MyServer.class, DefaultCtorNeutralizer.class);
        useInstrumenter(MyServerParent.class, DefaultCtorNeutralizer.class);

        System.out.println("Good scenario:");
        new MyServer();
    }
} 
